# 训练多位数字识别模型
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import os
import matplotlib as mpl

plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial']
mpl.rcParams['axes.unicode_minus'] = False
def generate_multi_digit_image(num_str, img_size=(128, 64), font_size=40):
    """生成多位数手写风格的图像"""
    img = Image.new('L', img_size, color=255)  # 白色背景
    draw = ImageDraw.Draw(img)
    
    # 加载手写体字体（需替换为本地路径）
    try:
        font = ImageFont.truetype("arial.ttf", font_size)  # 英文手写体
    except:
        font = ImageFont.load_default()  # 默认字体（备用）
    
    # 计算文本位置（居中）
    # 使用 textbbox 获取文本包围盒（相对于起始位置 (0,0)）
    bbox = draw.textbbox((0, 0), num_str, font=font)  # 返回 (x0, y0, x1, y1)
    text_width = bbox[2] - bbox[0]    # 文本宽度 = 右下角x - 左上角x
    text_height = bbox[3] - bbox[1]   # 文本高度 = 右下角y - 左上角y
    x = (img_size[0] - text_width) // 2
    y = (img_size[1] - text_height) // 2
    
    # 添加随机偏移（模拟手写不整齐）
    x += np.random.randint(-5, 5)
    y += np.random.randint(-5, 5)
    
    # 绘制文本（颜色随机深浅）
    text_color = 255 - np.random.randint(0, 100)  # 灰色到黑色
    draw.text((x, y), num_str, fill=text_color, font=font)
    
    # 添加噪声（模拟扫描失真）
    img_array = np.array(img)
    noise = np.random.normal(0, 10, img_array.shape).astype(np.uint8)
    noisy_img = np.clip(img_array + noise, 0, 255)
    return noisy_img

# 生成数据集（1-4位，共10000张图）
dataset_dir = "multi_digit_dataset"
os.makedirs(dataset_dir, exist_ok=True)

for _ in range(10000):
    # 随机生成1-4位数字
    length = np.random.randint(1, 5)
    num = np.random.randint(10**(length-1), 10**length)  # 如length=2时，10-99
    num_str = str(num)
    
    # 生成图像并保存
    img = generate_multi_digit_image(num_str)
    img_path = os.path.join(dataset_dir, f"{num_str}_{_}.png")
    Image.fromarray(img).save(img_path)

# 生成标签文件（txt格式：路径 数值）
with open(os.path.join(dataset_dir, "labels.txt"), "w") as f:
    for _ in range(10000):
        num_str = np.random.choice([str(np.random.randint(10**(l-1), 10**l)) for l in range(1,5)])
        img_path = os.path.join(dataset_dir, f"{num_str}_{_}.png")
        f.write(f"{img_path} {int(num_str)}\n")

img = generate_multi_digit_image("42", img_size=(128, 64), font_size=40)
plt.imshow(img, cmap='gray')
plt.axis('off')
plt.title("生成的多位数图像")
plt.show()